import re
import os
import time
import nltk
from nltk.tokenize import word_tokenize, sent_tokenize  # 分词和分句工具
from nltk.corpus import stopwords  # 停用词库
from nltk.probability import FreqDist  # 词频统计
from nltk.stem import WordNetLemmatizer  # 词形还原
from rich.progress import Progress, TextColumn, BarColumn, TimeElapsedColumn  # 进度条显示
from data_processing_common import sanitize_filename  # 文件名清理函数

def summarize_text_content(text, text_inference):
    """生成文本内容的摘要
    
    参数：
        text: 需要摘要的文本内容
        text_inference: 文本推理模型
        
    返回：
        生成的摘要文本
    """
    # 构建摘要提示词模板
    prompt = f"""请为以下文本提供简洁准确的摘要，聚焦主要观点和关键细节。
摘要字数限制在150字以内。

原文: {text}

摘要:"""

    # 调用模型生成摘要
    response = text_inference.create_completion(prompt)
    summary = response['choices'][0]['text'].strip()
    return summary

def process_single_text_file(args, text_inference, silent=False, log_file=None):
    """处理单个文本文件生成元数据
    
    参数：
        args: (文件路径, 文本内容)元组
        text_inference: 文本推理模型
        silent: 是否静默模式(默认False)
        log_file: 日志文件路径(可选)
        
    返回：
        包含文件元数据的字典
    """
    file_path, text = args
    start_time = time.time()

    # 创建进度条
    with Progress(
        TextColumn("[进度描述]{task.description}"),
        BarColumn(),
        TimeElapsedColumn()
    ) as progress:
        task_id = progress.add_task(f"正在处理 {os.path.basename(file_path)}", total=1.0)
        # 生成元数据
        foldername, filename, description = generate_text_metadata(text, file_path, progress, task_id, text_inference)

    end_time = time.time()
    time_taken = end_time - start_time

    # 输出处理结果
    message = f"文件: {file_path}\n耗时: {time_taken:.2f}秒\n描述: {description}\n分类目录: {foldername}\n生成文件名: {filename}\n"
    if silent:
        if log_file:
            with open(log_file, 'a') as f:
                f.write(message + '\n')
    else:
        print(message)
    return {
        'file_path': file_path,
        'foldername': foldername,
        'filename': filename,
        'description': description
    }

def process_text_files(text_tuples, text_inference, silent=False, log_file=None):
    """批量处理文本文件
    
    参数：
        text_tuples: (文件路径, 文本内容)元组列表
        text_inference: 文本推理模型
        silent: 是否静默模式
        log_file: 日志文件路径
        
    返回：
        处理结果列表
    """
    results = []
    for args in text_tuples:
        data = process_single_text_file(args, text_inference, silent=silent, log_file=log_file)
        results.append(data)
    return results

def generate_text_metadata(input_text, file_path, progress, task_id, text_inference):
    """生成文本文件的元数据(分类目录、文件名和描述)
    
    参数：
        input_text: 输入文本内容
        file_path: 文件路径
        progress: 进度条对象
        task_id: 进度任务ID
        text_inference: 文本推理模型
        
    返回：
        (分类目录, 文件名, 描述)元组
    """
    # 总处理步骤数
    total_steps = 3

    # 步骤1: 生成描述
    description = summarize_text_content(input_text, text_inference)
    progress.update(task_id, advance=1 / total_steps)

    # 步骤2: 生成文件名
    filename_prompt =  f"""根据以下摘要，生成一个能准确反映文档核心内容的文件名。
文件名限制在3个词以内。使用名词，避免使用"描述"、"展示"等动词开头。
不要包含"文本"、"文档"等类型词汇。只使用字母并用下划线连接单词。

摘要: {description}

示例:
1. 摘要: 一篇关于红军长征的文章
   文件名: 红军长征

2. 摘要: 讨论气候变化对北极熊影响的文章
   文件名: 气候变化_北极熊

请生成文件名。

只需输出文件名，不要包含其他文本。

文件名:"""
    filename_response = text_inference.create_completion(filename_prompt)
    filename = filename_response['choices'][0]['text'].strip()
    # 移除可能的前缀
    filename = re.sub(r'^文件名:\s*', '', filename, flags=re.IGNORECASE).strip()
    progress.update(task_id, advance=1 / total_steps)

    # 步骤3: 生成分类目录
    foldername_prompt = f"""根据以下摘要，生成最能代表该文档主题的通用分类。
这将用作分类目录名。分类名限制在2个词以内。使用名词，避免使用动词。
不要包含具体细节或文件名中的词汇，也不要使用"未分类"等通用词汇。

摘要: {description}

示例:
1. 摘要: 一篇关于红军长征的文章
   分类: 历史学

2. 摘要: 讨论气候变化对北极熊影响的文章
   分类: 环境

请生成分类。

只需输出分类，不要包含其他文本。

分类:"""
    foldername_response = text_inference.create_completion(foldername_prompt)
    foldername = foldername_response['choices'][0]['text'].strip()
    # 移除可能的前缀
    foldername = re.sub(r'^分类:\s*', '', foldername, flags=re.IGNORECASE).strip()
    progress.update(task_id, advance=1 / total_steps)

    # 定义不需要的词汇
    unwanted_words = set([
        '这个', '那个', '这些', '那些', '这里', '那里', '请', '注意', '附加', '笔记', 
        '文件夹', '名称', '确保', '是', '的', '在', '和', '基于', '生成', '文件', 
        '文档', '文本', '输出', '仅', '下面', '分类', '摘要', '关键', '细节', '信息',
        '主要', '想法', '概念', '关于', '通过', '为了', '从', '一个', '作为', '在',
        '我', '我们', '你', '他们', '他', '她', '它', '那', '哪个', '是', '有',
        '做', '但', '如果', '或', '因为', '每个', '一些', '没有', '不', '只有',
        '相同', '所以', '可以', '将', '应该', '现在', '新', '描述', '展示', '显示',
        '说明', '呈现', '提供', '讨论', '演示'
    ])
    stop_words = set(stopwords.words('english'))  # 英文停用词
    all_unwanted_words = unwanted_words.union(stop_words)
    lemmatizer = WordNetLemmatizer()  # 词形还原器

    def clean_ai_output(text, max_words):
        """清理AI生成的输出文本"""
        # 移除特殊字符和数字
        text = re.sub(r'[^\w\s]', ' ', text)
        text = re.sub(r'\d+', '', text)
        text = text.strip()
        # 处理驼峰命名
        text = re.sub(r'([a-z])([A-Z])', r'\1 \2', text)
        # 分词和词形还原
        words = word_tokenize(text)
        words = [word.lower() for word in words if word.isalpha()]
        words = [lemmatizer.lemmatize(word) for word in words]
        # 过滤不需要的词汇和重复词
        filtered_words = []
        seen = set()
        for word in words:
            if word not in all_unwanted_words and word not in seen:
                filtered_words.append(word)
                seen.add(word)
        # 限制单词数量
        filtered_words = filtered_words[:max_words]
        return '_'.join(filtered_words)

    # 处理文件名
    filename = clean_ai_output(filename, max_words=3)
    if not filename or filename.lower() in ('untitled', ''):
        # 从描述中提取关键词
        filename = clean_ai_output(description, max_words=3)
    if not filename:
        # 使用原始文件名
        filename = '文档_' + os.path.splitext(os.path.basename(file_path))[0]

    sanitized_filename = sanitize_filename(filename, max_words=3)

    # 处理分类目录名
    foldername = clean_ai_output(foldername, max_words=2)
    if not foldername or foldername.lower() in ('untitled', ''):
        # 从描述中提取关键词
        foldername = clean_ai_output(description, max_words=2)
        if not foldername:
            foldername = '文档集'

    sanitized_foldername = sanitize_filename(foldername, max_words=2)

    return sanitized_foldername, sanitized_filename, description